from datetime import date
from dateutil.relativedelta import relativedelta
from rqalpha.data.base_data_source import BaseDataSource
import pandas as pd
import os
import time
from rqalpha.core.execution_context import ExecutionContext
from rqalpha.const import EXECUTION_PHASE
from rqalpha.environment import Environment
from rqalpha.model.tick import TickObject
from rqalpha.interface import AbstractPriceBoard
from rqalpha.data.bar_dict_price_board import BarDictPriceBoard
from rqalpha.core.events import EVENT, Event
from rqalpha.my_factors.reform_data import get_source_tick_data, DataLoader, get_source_tick_data_from_merge_path,\
read_adj_factor_price
from rqalpha.const import TRADING_CALENDAR_TYPE
pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', None)


def convert_date_to_int(dt):
    t = dt.year * 10000 + dt.month * 100 + dt.day
    return t


def convert_dt_to_int(dt):
    t = dt.hour * 10000000 + dt.minute * 100000 + dt.second * 1000
    return t


class MyTickPriceBoard(BarDictPriceBoard):

    def __init__(self, *args, **kwargs):
        super(MyTickPriceBoard, self).__init__(*args, **kwargs)
        self._env = Environment.get_instance()
        # 这里没有绑定EVENT.tick是为了提前把数据加载好，再处理tick事件
        self._env.event_bus.add_listener(EVENT.PRE_TICK, self._on_tick)
        self.current_tick = None

    def get_a1(self, order_book_id):
        return self.current_tick._tick_dict['a1']

    def get_b1(self, order_book_id):
        return self.current_tick._tick_dict['b1']

    def get_last_price(self, order_book_id):
        # print('get last price---------------')
        # print(self.current_tick)
        if ExecutionContext.phase() == EXECUTION_PHASE.OPEN_AUCTION:
            return super(MyTickPriceBoard, self).get_last_price(order_book_id)
        if self._env.config.base.frequency == '1d':
            return super(MyTickPriceBoard, self).get_last_price(order_book_id)
        else:
            if self.current_tick is None:
                return 0
            return self.current_tick._tick_dict['last']

    def get_limit_down(self, order_book_id):
        if self._env.config.base.frequency == '1d':
            return super(MyTickPriceBoard, self).get_limit_down(order_book_id)
        else:
            return self.current_tick._tick_dict['limit_down']

    def get_limit_up(self, order_book_id):
        if self._env.config.base.frequency == '1d':
            return super(MyTickPriceBoard, self).get_limit_up(order_book_id)
        else:
            return self.current_tick._tick_dict['limit_up']

    def _get_last(self, order_book_id):
        # print('----------get_last')
        return self.current_tick._tick_dict['prev_close']

    def _on_tick(self, event):
        # todo 需要改
        # print('++++get_event+++++++:', event)
        self.current_tick = event.tick
        # print(event.tick)


class MyDataSource(BaseDataSource):
    fields_mapping = {'ask_vols': ['a1_v', 'a2_v', 'a3_v', 'a4_v', 'a5_v', 'a6_v', 'a7_v', 'a8_v', 'a9_v', 'a10_v'],
                      'asks': ['a1', 'a2', 'a3', 'a4', 'a5', 'a6', 'a7', 'a8', 'a9', 'a10'],
                      'bid_vols': ['b1_v', 'b2_v', 'b3_v', 'b4_v', 'b5_v', 'b6_v', 'b7_v', 'b8_v', 'b9_v', 'b10_v'],
                      'bids': ['b1', 'b2', 'b3', 'b4', 'b5', 'b6', 'b7', 'b8', 'b9', 'b10']}

    def __init__(self, path):
        super(MyDataSource, self).__init__(path, None)
        self._env = Environment.get_instance()
        run_method = self._env.config.base.run_method
        sel_type = self._env.config.base.select_type
        sys_freq = self._env.config.base.sys_freq
        sys_frequency = Environment.get_instance().config.base.frequency
        print(sys_frequency)
        self.sys_frequency = sys_frequency
        if sys_frequency in ['tick']:
            if run_method == 'horizon':
                symbols = self._env.config.base.all_symbols
                start_date = self._env.config.base.start_date
                end_date = self._env.config.base.end_date
                data_path = self._env.config.base.source_data_path
                merge_data_path = self._env.config.base.merge_data_path
                # print(merge_data_path)
                # time.sleep(500)
                if merge_data_path is not None:
                    self.ticks_data = DataLoader(sel_type, start_date, end_date, data_path, sys_freq,
                                                 symbols, merge_data_path).load_symbols_data_from_merge_path()
                else:
                    self.ticks_data = DataLoader(sel_type, start_date, end_date, data_path, sys_freq,
                                                 symbols, merge_data_path).load_symbols_data()
            if run_method == 'vertical':
                if self._env.config.base.merge_file_li is None:
                    source_file_li = self._env.config.base.source_file_li
                    book_id = self._env.config.base.symbol
                    trans_file_li = self._env.config.base.trans_file_li
                    self.ticks_data = get_source_tick_data(sel_type, source_file_li, trans_file_li, book_id, sys_freq)
                else:
                    merge_file_li = self._env.config.base.merge_file_li
                    self.ticks_data = get_source_tick_data_from_merge_path(merge_file_li)
            print(len(self.ticks_data))
            # print('+++++++++++++++++++++++++++++++++++++')
            # print(self.ticks_data.head(5))
            # print(self.ticks_data.shape, self.ticks_data['order_book_id'].unique().tolist(),
            #       self.ticks_data['trading_date'].unique().tolist())
            # time.sleep(500)

    def get_bar(self, instrument, dt, frequency):
        # # todo 需要改
        # print('get bar-------------')
        if frequency == '1d':
            # print('1d-----------')
            return super(MyDataSource, self).get_bar(instrument, dt, frequency)
        elif frequency == 'tick':
            print('tick-----------')

    def get_open_auction_bar(self, order_book_id, dt):
        obj = super(MyDataSource, self).get_open_auction_bar(order_book_id, dt)
        # dt_str = pd.to_datetime(dt, format='%Y-%m-%d').strftime('%Y-%m-%d')
        # print(order_book_id.order_book_id)
        # price = read_adj_factor_price(dt_str, order_book_id.order_book_id)
        # if price is not None:
        #     if obj['open'] != price:
        #         obj['open'] = price
        #     if obj['last'] != price:
        #         obj['last'] = price
        # print('get_open_auction_bar from my data_source:', obj['open'])
        return obj

    def _get_tick(self, instrument, short_df):
        for i, item in short_df.iterrows():
            date_int = convert_date_to_int(item['datetime'])
            dt_int = convert_dt_to_int(item['datetime'])
            date_dict = {'date': date_int, 'time': dt_int}
            field_dict = {}
            for key, field in self.fields_mapping.items():
                field_dict[key] = item[field].values.tolist()
            item_dict = item.to_dict()
            snapshot_dict = {**date_dict, **field_dict, **item_dict}
            # print(snapshot_dict)
            snapshot_dict['datetime'] = snapshot_dict['datetime'].to_pydatetime()
            yield TickObject(instrument, snapshot_dict)

    def get_merge_ticks(self, order_book_id_list, trading_date, last_dt=None):
        # 需要根据日期筛选这一天的所有数据
        # print(order_book_id_list, trading_date, last_dt)
        # print(order_book_id_list)
        if isinstance(order_book_id_list, set):
            for book in order_book_id_list:
                instrument = Environment.get_instance().get_instrument(book)
                order_book_id = instrument.order_book_id
                str_date = trading_date.strftime('%Y-%m-%d')
                # print('instruments:', instrument)
                # print(order_book_id, str_date)
                mask = (self.ticks_data['order_book_id'] == order_book_id) & \
                       (self.ticks_data['trading_date'] == str_date)
                short_df = self.ticks_data.loc[mask]
                # print(short_df)
                yield from self._get_tick(instrument, short_df)

    def available_data_range(self, frequency):
        return date(2005, 1, 1), date.today() - relativedelta(days=1)

    def get_trading_calendars(self):
        if self.sys_frequency == '1d':
            return super(MyDataSource, self).get_trading_calendars()
        else:
            sel_type = self._env.config.base.select_type
            if sel_type in ['d_currency']:
                trade_calender = pd.date_range('2005-01-01', '2023-12-31')
                return {TRADING_CALENDAR_TYPE.EXCHANGE: trade_calender}
            else:
                return super(MyDataSource, self).get_trading_calendars()
